/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 * Description:
 * Author: yanzhenghang
 * Create: 2024-09-06
 */

#include "transport_roce_mem.h"
#include "log.h"
#include "adapter_hal.h"
#include "adapter_hccp.h"
#include "adapter_rts.h"
#include "network_manager_pub.h"

namespace hccl {
using namespace std;
constexpr s32 REG_VALID = 1;
constexpr u32 WAIT_LINK_BUILD_DELAY_TIME_US = 10;
constexpr s32 QP_FLAG_RC = 0;          // flag: 0 = RC, 1= UD，其它预留
constexpr s32 OPBASE_QP_MODE_EXT = 4;  // 单算子模式(910B/910_93)的QP
constexpr u32 EIGHE_BIT = 8;
std::atomic<uint64_t> TransportRoceMem::sendWrHandle{0};
TransportRoceMem::TransportRoceMem(const std::unique_ptr<NotifyPool> &notifyPool, const HcclNetDevCtx &netDevCtx,
    const HcclDispatcher &dispatcher, u32 localRankId, u32 remoteRankId)
    : TransportMem(notifyPool, netDevCtx, dispatcher, localRankId, remoteRankId)
{}

TransportRoceMem::~TransportRoceMem()
{
    //  de rdmaSignal and Mr
    if (rdmaSignalMrHandle_ != nullptr) {
        HcclResult ret = HCCL_SUCCESS;
        ret = hrtRaDeRegGlobalMr(nicRdmaHandle_, rdmaSignalMrHandle_);
        if (ret != 0) {
            HCCL_ERROR("deReg rdmaSignal GlobalMr failed, ret[%d]", ret);
        }
    }
    // destory notify mem and notifyMem Mr
    if (notifyValueMemMrHandle_ != nullptr) {
        HcclResult ret = HCCL_SUCCESS;
        ret = hrtRaDeRegGlobalMr(nicRdmaHandle_, notifyValueMemMrHandle_);
        if (ret != 0) {
            HCCL_ERROR("deReg notify Mem Mr failed, ret[%d]", ret);
        }
    }
    if (notifyMem_.ptr() != nullptr) {
        notifyMem_.free();
    }
    // destory QP
    DestroyCqAndQp();
}

HcclResult TransportRoceMem::SetSocket(const std::shared_ptr<HcclSocket> &socket)
{
    CHK_SMART_PTR_NULL(socket);
    if (socket->GetStatus() != HcclSocketStatus::SOCKET_OK) {
        HCCL_ERROR("sockets does not connected");
        return HCCL_E_PARA;
    }
    socket_ = socket;
    return HCCL_SUCCESS;
}

HcclResult TransportRoceMem::GetRdmaHandle()
{
    RaResourceInfo raResourceInfo;
    CHK_RET(NetworkManager::GetInstance(deviceLogicId_).GetRaResourceInfo(raResourceInfo));
    auto it = raResourceInfo.nicSocketMap.find(socket_->GetLocalIp());
    if (it == raResourceInfo.nicSocketMap.end()) {
        HCCL_ERROR("[TransportRoceMem][GetRdmaHandle]nic socket handle did not found");
        return HCCL_E_PARA;
    }
    nicRdmaHandle_ = it->second.nicRdmaHandle;
    CHK_PTR_NULL(nicRdmaHandle_);
    HCCL_INFO(
        "TransportRoceMem GetNetworkResource deviceLogicId_[%d] nicRdmaHandle_[%p]", deviceLogicId_, nicRdmaHandle_);
    return HCCL_SUCCESS;
}

HcclResult TransportRoceMem::Connect(s32 timeoutSec)
{
    deviceLogicId_ = netDevCtx_->GetLogicId();
    CHK_PRT_RET(deviceLogicId_ == HOST_DEVICE_ID, HCCL_ERROR("deviceLogicId is same as host device id"), HCCL_E_INTERNAL);
    CHK_PTR_NULL(dispatcher_);
    CHK_SMART_PTR_NULL(notifyPool_);
    CHK_RET(notifyPool_->RegisterOp(socket_->GetTag()));
    CHK_RET(GetRdmaHandle());
    CHK_RET(CreateCqAndQp());
    CHK_RET(CreatSignalMesg());
    CHK_RET(CreateNotifyValueBuffer());
    CHK_RET(ExchangeNotifyValueBuffer());
    CHK_RET(QpConnect());
    CHK_RET(WaitQPLinkComplete(timeoutSec));
    CHK_RET(notifyPool_->UnregisterOp(socket_->GetTag()));
    return HCCL_SUCCESS;
}

HcclResult TransportRoceMem::TransportRdmaWithType(
    const RmaBufferSlice &localRmaBufferSlice, const RmaBufferSlice &remoteRmaBufferSlice, const rtStream_t &stream, const RdmaOp &rdmaOp)
{
    u64 processedOffset = 0;
    u64 remainingBytes = remoteRmaBufferSlice.len;
    u64 byteSizeChunk = 0;
    uint64_t localStartAddr = 0;
    uint64_t remoteStartAddr = 0;
    while (remainingBytes > 0) {
        localStartAddr = reinterpret_cast<uint64_t>(static_cast<u8 *>(localRmaBufferSlice.addr) + processedOffset);
        remoteStartAddr = reinterpret_cast<uint64_t>(static_cast<u8 *>(remoteRmaBufferSlice.addr) + processedOffset);
        byteSizeChunk = remainingBytes > MAX_RDMA_WQE_SIZE ? MAX_RDMA_WQE_SIZE : remainingBytes;
        std::shared_ptr<RemoteRdmaRmaBuffer> remoteRdmaRmaBuffer = dynamic_pointer_cast<RemoteRdmaRmaBuffer>(remoteRmaBufferSlice.rmaBuffer);
        std::shared_ptr<LocalRdmaRmaBuffer> localRdmaRmaBuffer = dynamic_pointer_cast<LocalRdmaRmaBuffer>(localRmaBufferSlice.rmaBuffer);
        struct sg_list list = {};
        list.addr = localStartAddr;
        list.len = byteSizeChunk;
        list.lkey = localRdmaRmaBuffer->GetKey();
        struct send_wr_v2 wr {};
        wr.wr_id = sendWrHandle.fetch_add(1);
        wr.buf_list = &list;
        wr.buf_num = 1;
        wr.dst_addr = remoteStartAddr;
        wr.rkey = remoteRdmaRmaBuffer->GetKey();
        wr.op = static_cast<u32>(rdmaOp);
        wr.send_flag = RA_SEND_SIGNALED;
        struct send_wr_rsp opRsp = {0};
        HCCL_DEBUG("Op type[%d], wr.wr_id[%llu], src addr[%p], dest addr[%p], len[%u]",
            rdmaOp,
            wr.wr_id,
            localRmaBufferSlice.addr,
            remoteRmaBufferSlice.addr,
            wr.buf_list->len);
        CHK_RET(HrtRaSendWrV2(dataQpInfo_.qpHandle, &wr, &opRsp, HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE));
        CHK_RET(DoorBellSend(dataQpInfo_.qpMode, wr, opRsp, stream));
        remainingBytes -= byteSizeChunk;
        processedOffset += byteSizeChunk;
    }
    return HCCL_SUCCESS;
}

HcclResult TransportRoceMem::TransportIpc(
    const RmaBufferSlice &dstRmaBufferSlice, const RmaBufferSlice &srcRmaBufferSlice, const rtStream_t &stream)
{
    Stream hcclStream(stream);
    CHK_RET(dispatcher_->MemcpyAsync(dstRmaBufferSlice.addr, dstRmaBufferSlice.len, srcRmaBufferSlice.addr,
        srcRmaBufferSlice.len, HcclRtMemcpyKind::HCCL_RT_MEMCPY_KIND_DEVICE_TO_DEVICE, hcclStream, remoteRankId_,
        hccl::LinkType::LINK_HCCS));
    return HCCL_SUCCESS;
}

HcclResult TransportRoceMem::Write(const RmaBufferSlice &remoteRmaBufferSlice, const RmaBufferSlice &localRmaBufferSlice, const rtStream_t &stream)
{
    CHK_RET(TransportRdmaWithType(localRmaBufferSlice, remoteRmaBufferSlice, stream, RdmaOp::OP_WRITE));
    return HCCL_SUCCESS;
}

HcclResult TransportRoceMem::Read(const RmaBufferSlice &localRmaBufferSlice, const RmaBufferSlice &remoteRmaBufferSlice, const rtStream_t &stream)
{
    CHK_RET(TransportRdmaWithType(localRmaBufferSlice, remoteRmaBufferSlice, stream, RdmaOp::OP_READ));
    return HCCL_SUCCESS;
}

HcclResult TransportRoceMem::AddOpFence(const rtStream_t &stream)
{
    auto opType = MemType::SEND_NOTIFY_MEM;
    struct sg_list list = {};
    list.addr = reinterpret_cast<uint64_t>(rdmaSignal_[0].addr);
    list.len = notifyMemMsg_[static_cast<u32>(opType)].len;
    list.lkey = rdmaSignal_[0].lkey;

    struct send_wr_v2 wr {};
    wr.wr_id = sendWrHandle.fetch_add(1);
    wr.buf_list = &list;
    wr.buf_num = 1; /* 此处list只有一个，设置为1 */
    wr.dst_addr = static_cast<u64>(reinterpret_cast<uintptr_t>(notifyMemMsg_[static_cast<u32>(opType)].addr));
    wr.rkey = notifyMemMsg_[static_cast<u32>(opType)].rkey;
    wr.op = static_cast<u32>(RdmaOp::OP_READ); /* RDMA_READ: 4 */
    wr.send_flag = RA_SEND_SIGNALED | RA_SEND_FENCE;
    struct send_wr_rsp opRsp = {};
    CHK_RET(HrtRaSendWrV2(dataQpInfo_.qpHandle, &wr, &opRsp, HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE));
    CHK_RET(DoorBellSend(dataQpInfo_.qpMode, wr, opRsp, stream));
    hccl::Stream hcclStream(stream);
    CHK_RET(LocalIpcNotify::Wait(hcclStream,
        dispatcher_,
        remoteIsendDoneSignal_,
        INVALID_VALUE_STAGE,
        NOTIFY_INVALID_WAIT_TIME,
        localRankId_,
        remoteRankId_));
    HCCL_DEBUG("AddOpFence wr.wr_id[%llu], src addr[%p], dest addr[%p], len[%u], streamid[%d]",
        wr.wr_id,
        rdmaSignal_[0].addr,
        notifyMemMsg_[static_cast<u32>(opType)].addr,
        wr.buf_list->len,
        hcclStream.id());
    return HCCL_SUCCESS;
}

HcclResult TransportRoceMem::CreateCqAndQp()
{
    dataQpInfo_.flag = QP_FLAG_RC;
    dataQpInfo_.qpMode = OPBASE_QP_MODE_EXT;
    CHK_RET(CreateQpWithCq(nicRdmaHandle_, -1, -1, nullptr, nullptr, dataQpInfo_, true, true));
    return HCCL_SUCCESS;
}

HcclResult TransportRoceMem::QpConnect()
{
    CHK_RET(HrtRaQpConnectAsync(dataQpInfo_.qpHandle, socket_->GetFdHandle()));

    return HCCL_SUCCESS;
}

HcclResult TransportRoceMem::RecoverNotifyMsg(MemMsg *remoteRdmaSignal, u64 signalNum)
{
    if (signalNum <= 0) {
        return HCCL_E_NOT_FOUND;
    }
    MemType tmpMemType = MemType::MEM_TYPE_RESERVED;
    for (u64 i = 0; i < signalNum; i++) {
        HCCL_DEBUG("recv mrRegFlag:[%d] notifyAddr:[%p] len:[%lu] memType:[%d], rkey:[%u]  ",
            (remoteRdmaSignal + i)->mrRegFlag,
            (remoteRdmaSignal + i)->addr,
            (remoteRdmaSignal + i)->len,
            static_cast<int>((remoteRdmaSignal + i)->memType),
            (remoteRdmaSignal + i)->lkey);
        tmpMemType = (remoteRdmaSignal + i)->memType;
        if ((remoteRdmaSignal + i)->memType == MemType::NOTIFY_SRC_MEM) {
            tmpMemType = MemType::SEND_NOTIFY_MEM;
            notifyMemMsg_[tmpMemType].mrRegFlag = (remoteRdmaSignal + i)->mrRegFlag;
            notifyMemMsg_[tmpMemType].addr = (remoteRdmaSignal + i)->addr;
            notifyMemMsg_[tmpMemType].len = (remoteRdmaSignal + i)->len;
            notifyMemMsg_[tmpMemType].memType = MemType::SEND_NOTIFY_MEM;
            notifyMemMsg_[tmpMemType].rkey = (remoteRdmaSignal + i)->lkey;
        }
    }

    return HCCL_SUCCESS;
}

HcclResult TransportRoceMem::CreatSignalMesg()
{
    CHK_RET(GetNotifySize());
    CHK_RET(CreateRdmaSignal(remoteIsendDoneSignal_, rdmaSignal_[0], MemType::RECV_NOTIFY_MEM));
    return HCCL_SUCCESS;
}

HcclResult TransportRoceMem::GetNotifySize()
{
    DevType devType;
    CHK_RET(hrtHalGetDeviceType(deviceLogicId_, devType));

    if (devType == DevType::DEV_TYPE_910) {
        notifySize_ = 8;  // 910A 每个notify占8个字节
    } else if ((devType == DevType::DEV_TYPE_910B) || (devType == DevType::DEV_TYPE_910_93)) {
        notifySize_ = 4;  // 910B/910_93 每个notify占4个字节
    } else {
        notifySize_ = 8;  // 其余芯片类型每个notify占8个字节
    }
    HCCL_INFO("devType[%d] notifySize[%d]", devType, notifySize_);
    return HCCL_SUCCESS;
}

HcclResult TransportRoceMem::ExchangeNotifyValueBuffer()
{
    CHK_RET(socket_->Send(
        &notifyMemMsg_[static_cast<u32>(MemType::NOTIFY_SRC_MEM)], sizeof(MemMsg) * REMOTE_RDMA_SIGNAL_SIZE));
    HCCL_DEBUG("send mrRegFlag:[%d] notifyAddr:[%p] len:[%lu] memType:[%d], rkey:[%u]  ",
        notifyMemMsg_[static_cast<u32>(MemType::NOTIFY_SRC_MEM)].mrRegFlag,
        notifyMemMsg_[static_cast<u32>(MemType::NOTIFY_SRC_MEM)].addr,
        notifyMemMsg_[static_cast<u32>(MemType::NOTIFY_SRC_MEM)].len,
        notifyMemMsg_[static_cast<u32>(MemType::NOTIFY_SRC_MEM)].memType,
        notifyMemMsg_[static_cast<u32>(MemType::NOTIFY_SRC_MEM)].lkey);
    MemMsg remoteNotifyValue[REMOTE_RDMA_SIGNAL_SIZE];
    CHK_RET(socket_->Recv(remoteNotifyValue, sizeof(MemMsg) * REMOTE_RDMA_SIGNAL_SIZE));
    CHK_RET(RecoverNotifyMsg(remoteNotifyValue, REMOTE_RDMA_SIGNAL_SIZE));
    return HCCL_SUCCESS;
}

HcclResult TransportRoceMem::CreateRdmaSignal(
    std::shared_ptr<LocalIpcNotify> &localNotify, MemMsg &rdmaSignalInfo, MemType notifyType)
{
    u64 notifyOffset = 0;
    u64 notifyBaseVa = 0;  // notify寄存器虚拟地址
    u64 notifyTotalSize = 0;

    devicePhyId_ = netDevCtx_->GetPhyId();
    RemoteRankInfo info(devicePhyId_, remoteRankId_);
    CHK_RET(SalGetBareTgid(reinterpret_cast<u32 *>(&info.remotePid)));  // 当前进程id
    CHK_RET(notifyPool_->Alloc(socket_->GetTag(), info, localNotify));
    // 设置remote id
    s64 recvId = 0xFFFFFFFF00000000 | (static_cast<s64>(info.remotePid) & 0xFFFFFFFF);
    CHK_RET(localNotify->Grant(recvId));

    CHK_RET(HrtRaGetNotifyBaseAddr(nicRdmaHandle_, &notifyBaseVa, &notifyTotalSize));
    CHK_RET(localNotify->GetNotifyOffset(notifyOffset));
    u64 notifyVa = notifyBaseVa + notifyOffset;
    rdmaSignalInfo.mrRegFlag = 0;
    rdmaSignalInfo.addr = reinterpret_cast<void *>(static_cast<uintptr_t>(notifyVa));
    rdmaSignalInfo.len = notifySize_;
    rdmaSignalInfo.memType = notifyType;

    struct mr_info mrInfo = {};
    mrInfo.addr = rdmaSignalInfo.addr;
    mrInfo.size = rdmaSignalInfo.len;
    mrInfo.access = access_;
    CHK_RET(hrtRaRegGlobalMr(nicRdmaHandle_, mrInfo, rdmaSignalMrHandle_));
    rdmaSignalInfo.lkey = mrInfo.lkey;

    HcclSignalInfo notifyInfo{INVALID_U64};
    CHK_RET(localNotify->GetNotifyData(notifyInfo));
    HCCL_INFO("CreateRdmaSignal localNotify id[%llu]", notifyInfo.resId);
    return HCCL_SUCCESS;
}

HcclResult TransportRoceMem::CreateNotifyValueBuffer()
{
    if (notifyMem_.ptr() == nullptr) {
        u64 notifyVaule = 1;  // notify值写1表示record
        notifyMem_ = DeviceMem::alloc(notifyValueSize_);
        CHK_PRT_RET(
            !notifyMem_, HCCL_ERROR("[Create][NotifyValueBuffer]In lbv notify buffer, malloc failed."), HCCL_E_MEMORY);
        HCCL_DEBUG("create notify value buffer[%p], size[%u]", notifyMem_.ptr(), notifySize_);

        CHK_RET(hrtMemSyncCopy(notifyMem_.ptr(),
            notifyMem_.size(),
            &notifyVaule,
            notifySize_,
            HcclRtMemcpyKind::HCCL_RT_MEMCPY_KIND_HOST_TO_DEVICE));
    }

    struct mr_info mrInfo = {nullptr};
    mrInfo.addr = notifyMem_.ptr();
    mrInfo.size = notifySize_;
    mrInfo.access = access_;
    CHK_RET(hrtRaRegGlobalMr(nicRdmaHandle_, mrInfo, notifyValueMemMrHandle_));
    notifyMemMsg_[static_cast<u32>(MemType::NOTIFY_SRC_MEM)].mrRegFlag = REG_VALID;
    notifyMemMsg_[static_cast<u32>(MemType::NOTIFY_SRC_MEM)].addr = notifyMem_.ptr();
    notifyMemMsg_[static_cast<u32>(MemType::NOTIFY_SRC_MEM)].len = notifySize_;
    notifyMemMsg_[static_cast<u32>(MemType::NOTIFY_SRC_MEM)].memType = MemType::NOTIFY_SRC_MEM;
    notifyMemMsg_[static_cast<u32>(MemType::NOTIFY_SRC_MEM)].lkey = mrInfo.lkey;
    HCCL_DEBUG("notifyValueMem_=%p", notifyMem_.ptr());
    return HCCL_SUCCESS;
}

HcclResult TransportRoceMem::DoorBellSend(
    const s32 qpMode, const send_wr_v2 &sendWrV2, const send_wr_rsp &opRsp, void *stream)
{
    struct send_wr sendwr;
    sendwr.buf_list = sendWrV2.buf_list;
    sendwr.buf_num = sendWrV2.buf_num;
    sendwr.dst_addr = sendWrV2.dst_addr;
    sendwr.rkey = sendWrV2.rkey;
    sendwr.op = sendWrV2.op;
    sendwr.send_flag = sendWrV2.send_flag;
    u32 dbIndex = static_cast<u32>(opRsp.db.db_index);
    u64 dbInfo = static_cast<u64>(opRsp.db.db_info);
    CHK_RET(RdmaDbSend(dbIndex, dbInfo, sendwr, stream));
    return HCCL_SUCCESS;
}

HcclResult TransportRoceMem::RdmaDbSend(u32 dbindex, u64 dbinfo, const struct send_wr &sendWr, rtStream_t stream)
{
    hccl::Stream hcclStream(stream);
    s32 ret = dispatcher_->RdmaSend(dbindex, dbinfo, sendWr, hcclStream, remoteRankId_);
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("[RdmaDbSend]errNo[0x%016llx] rdma db send fail, "
                   "return[%d]. para: dbindex[%u]dbinfo[%llu].",
            HCCL_ERROR_CODE(HCCL_E_INTERNAL),
            ret,
            dbindex,
            dbinfo),
        HCCL_E_INTERNAL);
    return HCCL_SUCCESS;
}

HcclResult TransportRoceMem::WaitQPLinkComplete(s32 timeoutSec)
{
    auto startTime = chrono::steady_clock::now();
    maxTimeOut_ = std::chrono::seconds(timeoutSec);
    while ((chrono::steady_clock::now() - startTime) < maxTimeOut_) {
        HcclResult ret = GetQpStatus();
        if (ret == HCCL_E_AGAIN) {
            SaluSleep(WAIT_LINK_BUILD_DELAY_TIME_US);
            continue;
        }
        if (ret == HCCL_SUCCESS) {
            HCCL_INFO("TransportRoceMem QP connect success");
        } else {
            HCCL_ERROR("TransportRoceMem QP connect failed, ret[%d]!", ret);
        }
        return ret;
    }
    HCCL_RUN_INFO(
        "WaitBuildLinkComplete timeOut[%d] s, localRank[%u], remoteRank[%u]", timeoutSec, localRankId_, remoteRankId_);
    return HCCL_E_TIMEOUT;
}

HcclResult TransportRoceMem::GetQpStatus()
{
    int qpStatus = 0;
    s32 ret = 0;

    ret = hrtGetRaQpStatus(dataQpInfo_.qpHandle, &qpStatus);
    if (ret != 0) {
        return HCCL_E_INTERNAL;
    } else if (ret == 0 && qpStatus != 1) {  // 为1时，qp 建链成功
        return HCCL_E_AGAIN;
    }
    return HCCL_SUCCESS;
}

HcclResult TransportRoceMem::DestroyCqAndQp()
{
    HCCL_INFO("TransportRoceMem DestroyCqAndQp");
    CHK_RET(DestroyQpWithCq(dataQpInfo_, true));
    return HCCL_SUCCESS;
}

}  // namespace hccl
